#  /**
#  * Copyright (c) 2022 Beijing Jiaotong University
#  * PhotLab is licensed under [Open Source License].
#  * You can use this software according to the terms and conditions of the [Open Source License].
#  * You may obtain a copy of [Open Source License] at: [https://open.source.license/]
#  *
#  * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
#  * EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
#  * MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
#  *
#  * See the [Open Source License] for more details.
#  */
#  /**
#  * Author: Zheng Wang
#  * Created: Jun. 27, 2022
#  * Supported by: National Key Research and Development Program of China
#  */

import logging
import phot
import inspect
import copy
import threading
from phot_engine.utils.unit_conversion import UnitConversion
from phot_engine.globals.engine_exception import ArithmeticUnitError
from phot_engine.computing_graph.vertex import Vertex
from phot_engine.globals.engine_globals import init_engine_globals
from phot_engine.globals.arithmetic_unit_status import ArithmeticUnitStatus
from phot_engine.globals.backend import Backend
from phot_engine.logical_structure.resource_manager import ResourceManager

# If engine allows library to draw figure
SHOW_FIGURE = False


def _check_conversion_necessity(old_unit, new_unit):
    """
        check conversion necessity.If old unit equals new unit,there is no necessity to convert unit
    :param old_unit:
    :param new_unit:
    :return:
    """
    if len(old_unit) > 0 and len(new_unit) > 0 and (old_unit != new_unit):
        return True
    else:
        return False


class ArithmeticUnit(threading.Thread):
    """
    The distribution ArithmeticUnit does the following:
    1. Generate a vertex parameter dictionary to automatically pass function parameter
    2. Generate a dictionary of vertex input signals
    3. Generate a dictionary of vertex output signals
    4. Modify edge_list for signal distribution
    """
    lock = threading.RLock()  # 线程锁

    def __init__(self, vertex, edge_list, error_queue, recommend_gpu):
        if isinstance(vertex, Vertex):
            super(ArithmeticUnit, self).__init__()
            self.__vertex = vertex
            self.__edge_list = edge_list
            self.__uc = UnitConversion()
            self.__global_parameter = None
            self.__error_queue = error_queue
            self.__phot_classes_name, self.__phot_functions_name = init_engine_globals(module_name=phot)
            self.__arithmetic_unit_status = ArithmeticUnitStatus.THREAD_INITIALIZED
            self.__resource_manager = ResourceManager()
            self.recommend_gpu = recommend_gpu
        else:
            raise TypeError("vertex should be Vertex type,not {}".format(type(vertex)))

    def get_arithmetic_unit_status(self):
        return self.__arithmetic_unit_status

    def get_edge_list(self):
        return self.__edge_list

    def get_vertex(self):
        return self.__vertex

    def get_vertex_id(self):
        return self.__vertex.get_vertex_id()

    def set_global_parameter(self, global_parameter):
        self.__global_parameter = global_parameter

    def __is_function_exist(self, function_name):
        """
            judging a function name if exists in phot function list or not
        :param function_name:
        :return:
        """
        if function_name in self.__phot_functions_name:
            return True
        return False

    def __is_class_exist(self, class_name):
        """
            judging a class name if exists in phot class list or not
        :param class_name:
        :return:
        """
        if class_name in self.__phot_classes_name:
            return True
        return False

    def __get_class(self, class_name):
        """
            获取底层运算库类
            Gets a class in the phot classes
        :param class_name:
        :return:
        """
        if self.__is_class_exist(class_name):
            return getattr(phot, class_name)
        else:
            raise ArithmeticUnitError(
                "Thread:{},class_name:{} is not not in {}"
                    .format(self.name, class_name, self.__phot_classes_name))

    def __get_class_function(self, class_name, function_name):
        """
            获取底层运算库类中的成员方法
            Gets a member method in the phot class
        :param class_name:
        :param function_name:
        :return:
        """
        cls = self.__get_class(class_name)
        func_pointer = getattr(cls, function_name)
        return func_pointer

    def __get_function(self, function_name):
        """
            通过函数名返回底层运算库函数指针
            Returns the phot module's function pointer by function name
        :param function_name:
        :return:
        """
        if self.__is_function_exist(function_name):
            logging.debug("Arithmetic unit thread:{},get a function from phot:{}"
                          .format(self.name,
                                  getattr(phot, function_name)))

            return getattr(phot, function_name)
        else:
            raise ArithmeticUnitError(
                "Thread:{},function_name:{} is not not in {}"
                    .format(self.name, function_name, self.__phot_functions_name))

    def __instantiate_phot_class(self, class_name, instantiation_parameter_dict):
        """
            实例化phot中的类
            Instantiate a class in phot
        :param class_name:class name
        :param instantiation_parameter_dict: init function parameter dict
        :return:
        """
        phot_class = self.__get_class(class_name)
        init_func_pointer = self.__get_class_function(class_name, "__init__")
        filtered_dict = self.__gen_function_param_dict(func_pointer=init_func_pointer,
                                                       param_dict=instantiation_parameter_dict)
        phot_obj = phot_class(**filtered_dict)
        return phot_obj

    def __call_phot_obj_function(self, phot_class_name, function_name, instantiation_parameter_dict,
                                 input_signal_list=None):
        """
            根据类名,方法名,实例化参数字典,输入信号列表,调用底层运算库函数
            According to the class name, method name, instantiation parameter dictionary, input signal dictionary,
            call the phot function
        :param phot_class_name:
        :param function_name:
        :param instantiation_parameter_dict:
        :param input_signal_list:
        :return:
        """
        try:
            phot_obj = self.__instantiate_phot_class(class_name=phot_class_name,
                                                     instantiation_parameter_dict=instantiation_parameter_dict)

            phot_obj_function_pointer = getattr(phot_obj, function_name)
            logging.debug("Arithmetic unit thread:{},object in phot:{},input signal:{},function in phot:{}"
                          .format(self.name, phot_obj, input_signal_list, phot_obj_function_pointer))
            if input_signal_list is None:
                return phot_obj_function_pointer()
            else:
                return phot_obj_function_pointer(*input_signal_list)
        except Exception as e:
            raise e

    def __gen_function_param_dict(self, func_pointer, param_dict):
        """
            Function parameter lists are automatically parsed by function pointer
            and the input dictionaries are filtered
        :param func_pointer:
        :param param_dict:
        :return:
        """
        if inspect.isfunction(func_pointer) or inspect.ismethod(func_pointer):
            filtered_dict = dict()
            parameter_list = [parameter.name for parameter in inspect.signature(func_pointer).parameters.values()]
            if param_dict is not None and isinstance(param_dict, dict):
                for key, value in param_dict.items():
                    if key in parameter_list:
                        filtered_dict.update({key: value})
                return filtered_dict
            else:
                return None
        else:
            raise ArithmeticUnitError("You should input a function pointer,not {}".format(type(func_pointer)))

    def __gen_input_signal_list(self):
        """
            生成输入信号列表,如果当前顶点入度为0则证明无需输入信号,返回空
            输入信号列表的生成方式为：
                1. 在传入的edge_list中寻找以当前图顶点为终点的边
                2. 根据当前边中存储的信号设置 value
            Generate input signal list, if the current vertex entry degree is 0,
            the proof that no input signal, return None
            The input signal tuple is generated as follows:
            1. Find the edge in the passed edge_list that ends at the vertex of the current graph
            2. Set value based on the signal stored in the current edge
        :return:
        """
        if 0 == self.__vertex.get_indegree():
            return None
        else:
            input_signal_list = []
            vertex_id = self.__vertex.get_vertex_id()
            # 对于临界资源edge_list,在查询时暂时不添加线程锁
            for edge in self.__edge_list:
                if edge.destination_port.vertex_id == vertex_id:
                    # 使用深拷贝构造输入信号列表,防止传入和历史信号相同的内存地址,导致出现修改历史计算结果的问题
                    input_signal_list.insert(int(edge.destination_port.port_id), copy.deepcopy(edge.signal))
            logging.debug(
                "The list of input signals constructed for vertex :{} is {},input signal list length:{}"
                    .format(self.__vertex.get_vertex_id(), input_signal_list, len(input_signal_list)))

            return input_signal_list

    # 当前版本未用到
    def __execute_phot_global_function(self, phot_global_parameter_dict):
        """
            execute phot global initialize function,if phot has the global parameter initialization
        :param phot_global_parameter_dict:
        :return:
        """
        return True

    def __execute_phot_init_function(self):
        """
            执行底层运算库初始化函数
        :return:
        """
        computing_method = self.__vertex.get_parallel_computing_method()
        logging.debug("{} will allocate on {}".format(self.__vertex.get_vertex_id(), computing_method))
        if computing_method == "GPU":
            try:
                import cupy
                phot.config(plot=SHOW_FIGURE, backend=Backend.CUPY.value, gpu_device=self.recommend_gpu[0])
                print(self.recommend_gpu[0], self.__vertex.get_vertex_id())
            except ModuleNotFoundError as import_e:
                logging.warning("Engine attempt to allocate cupy as calculation library, "
                                "but cupy doesn't install. Change to use torch.")
                try:
                    import torch
                    if torch.cuda.is_available() is False:
                        phot.config(plot=SHOW_FIGURE, backend=Backend.NUMPY.value)
                    else:
                        phot.config(plot=SHOW_FIGURE, backend=Backend.TORCH.value, gpu_device=self.recommend_gpu[0])
                except ModuleNotFoundError as import_e:
                    logging.warning("Engine attempt to allocate torch as calculation library, "
                                    "but torch doesn't install. Change to use numpy.")
                    phot.config(plot=SHOW_FIGURE, backend=Backend.NUMPY.value)
        else:
            phot.config(plot=SHOW_FIGURE, backend=Backend.NUMPY.value)

    def __execute_phot_function(self, function_pointer, function_param):
        """
            execute phot functions,function_param has two type.
            If it is list type,the parameter represents input signal,else is function parameter
        :param function_pointer:
        :param function_param:
        :return:
        """
        if isinstance(function_param, dict):
            # 只有参数
            filtered_param_dict = self.__gen_function_param_dict(function_pointer, function_param)
            return function_pointer(**filtered_param_dict)
        elif isinstance(function_param, list):
            # 只有信号
            return function_pointer(*function_param)
        elif isinstance(function_param, tuple):
            # 当信号和参数混合使用时
            logging.debug(function_pointer)
            logging.debug("function_pointer signal[0]: {}".format(function_param[0]))
            logging.debug("function_pointer parameter[1]: {}".format(function_param[1]))
            return function_pointer(*function_param[0], **function_param[1])

    def __scatter_output_signal_tuple(self, output_signal_tuple):
        """
            scatter output signal tuple for vertex in computing graph
            为图顶点分配输出信号
        :param output_signal_tuple:
        :return:
        """
        self.lock.acquire()
        logging.debug("Arithmetic Unit Thread:{},scatters output signal for vertex:{},"
                      "signal tuple id is:{},signal:{}"
                      .format(self.name, self.__vertex.vertex_id,
                              id(output_signal_tuple), output_signal_tuple))
        if not isinstance(output_signal_tuple, tuple):
            self.__alter_edge_list(0, output_signal_tuple)
        else:
            for port_id in range(0, len(output_signal_tuple)):
                self.__alter_edge_list(port_id, output_signal_tuple[port_id])
        self.lock.release()

    def __alter_edge_list(self, port_id, signal):
        """
            alter the signal in edge list
            修改边列表中的边信息
        :param port_id:
        :param signal:
        :return:
        """
        vertex_id = self.__vertex.get_vertex_id()
        logging.debug("Arithmetic Unit Thread:{} to alter edge list length:{}".format(self.name,
                                                                                      len(self.__edge_list)))
        for edge in self.__edge_list:
            if vertex_id == edge.source_port.vertex_id and port_id == int(edge.source_port.port_id):
                logging.debug("Arithmetic Unit Thread:{} altered edge list."
                              "Edge list id:{},scattered signal:{} for vertex:{}".format(self.name,
                                                                                         id(self.__edge_list),
                                                                                         signal,
                                                                                         self.__vertex.get_vertex_id()))
                edge.signal = signal

        # print("顶点:{} 输出信号分发完毕,仿真任务边列表:".format(vertex_id), self.__edge_list)

    def __gen_standard_dict(self, nonstandard_vertex_param_dict: dict):
        """

        :param nonstandard_vertex_param_dict:
        has value like that:
        {'length': param_name: length,param_value: 262144,param_unit: ,param_phot_unit:,
         'seed': param_name: seed,param_value: 1,param_unit: ,param_phot_unit:}
        we want is:
        {'length': 262144, 'seed': 1}
        :return:
        """
        res_dict = {}
        for param_key in nonstandard_vertex_param_dict:
            one_param = nonstandard_vertex_param_dict[param_key]
            res_dict[param_key] = one_param.param_value
        return res_dict

    def __determinate_phot_function_call_method(self, input_signal, filtered_std_param_dict, function_pointer):
        """
            根据输入信号和组件参数字典来决定底层运算库的函数调用方法
                1. 只有输入信号,没有组件参数
                2. 没有输入信号,只有组件参数
                3. 既有输入信号，又有组件参数
        :param input_signal:
        :param filtered_std_param_dict:
        :param function_pointer:
        :return:
        """
        logging.debug("""Determinate Vertex:{} 
                         phot function call method,input signal:{},type:{},
                         filtered std parameter dictionary:{},
                         function pointer:{}""".format(self.__vertex.get_vertex_id(),
                                                       input_signal,
                                                       type(input_signal),
                                                       filtered_std_param_dict,
                                                       function_pointer))
        if (input_signal is not None and len(input_signal) > 0) and (len(filtered_std_param_dict) == 0):
            # 存在输入信号
            self.__scatter_output_signal_tuple(
                self.__execute_phot_function(function_pointer, input_signal))
        elif len(filtered_std_param_dict) > 0 and (input_signal is None or len(input_signal) == 0):
            # 只存在组件参数
            self.__scatter_output_signal_tuple(
                self.__execute_phot_function(function_pointer, filtered_std_param_dict))
        elif len(filtered_std_param_dict) > 0 and input_signal is not None and len(input_signal) > 0:
            # 即存在组件参数也存在输入信号
            self.__scatter_output_signal_tuple(
                self.__execute_phot_function(function_pointer, (input_signal, filtered_std_param_dict)))
        else:
            raise ArithmeticUnitError("Error in __determinate_phot_function_call_method(), "
                                      "Detected a nonexistent way to call a phot function.")

    def __execute_normal_vertex_processing_logic(self):
        """
        对于正常图顶点
        计算引擎自动调用底层运算库逻辑代码,主要工作过程如下所示:
            对于正常图顶点
                1.寻找底层运算库类和函数
                2.器件参数分发
                3.输入信号分发
                4.函数调用
                5.输出信号调度
            对于空的图顶点
                1.根据输入边分发信号
        The computing engine automatically calls the logic code of the pyphot library
        and the main working process is as follows:
            For normal Vertex:
                1. Find the underlying library classes and functions
                2. Distribute device parameters
                3. Input signal distribution
                4. Function calls
                5. Output signal scheduling
            For None Vertex:
                1. scatter signal by inout edge
        :return:
        """
        logging.debug("Arithmetic Unit Thread:{} prepares to calculate vertex:{},class name:{},function name:{}"
                      .format(self.name,
                              self.__vertex.get_vertex_id(),
                              self.__vertex.phot_lib_class_name,
                              self.__vertex.phot_lib_func_name))

        vertex_param_dict = self.__vertex.get_param_dict()
        # You need to pick up only key and value from the vertex_param dict
        standard_param_dict = self.__gen_standard_dict(vertex_param_dict)
        class_name = self.__vertex.phot_lib_class_name
        function_name = self.__vertex.phot_lib_func_name

        input_signal_list = self.__gen_input_signal_list()
        logging.debug("Arithmetic Unit Thread:{},vertex:{},input signal:{}"
                      .format(self.name, self.__vertex.vertex_id, input_signal_list))

        if self.__execute_phot_global_function(None):
            logging.debug("Arithmetic Unit Thread:{} phot's global parameter:{}"
                          .format(self.name, self.__global_parameter))
            """
            底层运算库函数有两种:
                1.通过实例化相应对象调用
                2.直接通过函数调用
                    1) 其中一种函数是代表了一个器件,不是为它传参数，而是传信号
                    2) 还有一种函数,只是代表了一个函数,不是为它传信号,而是参数
            调用函数的过程需要上锁
            """
            if class_name == function_name:
                function_pointer = self.__get_function(function_name)
                filtered_std_param_dict = self.__gen_function_param_dict(function_pointer, standard_param_dict)
                self.__determinate_phot_function_call_method(input_signal=input_signal_list,
                                                             filtered_std_param_dict=filtered_std_param_dict,
                                                             function_pointer=function_pointer)
            else:
                self.__scatter_output_signal_tuple(
                    self.__call_phot_obj_function(phot_class_name=class_name,
                                                  function_name=function_name,
                                                  instantiation_parameter_dict=standard_param_dict,
                                                  input_signal_list=input_signal_list))

            self.__arithmetic_unit_status = ArithmeticUnitStatus.THREAD_ACCOMPLISHED
        else:
            err = ArithmeticUnitError("运算单元-底层运算库全局初始化失败")
            self.__arithmetic_unit_status = ArithmeticUnitStatus.THREAD_EXCEPTION
            self.__error_queue.put(err)

    def __execute_none_vertex_processing_logic(self):
        input_signal_list = self.__gen_input_signal_list()
        self.__scatter_output_signal_tuple((*input_signal_list,))
        self.__arithmetic_unit_status = ArithmeticUnitStatus.THREAD_ACCOMPLISHED
        logging.debug("This is a None Vertex:{},thread id:{}, Only scatter input signal."
                      .format(self.__vertex.get_vertex_id(),
                              self.name))

    def run(self) -> None:

        try:
            self.__arithmetic_unit_status = ArithmeticUnitStatus.THREAD_RUNNING
            logging.info("Start to calculate vertex:{},in Arithmetic Unit Thread:{},is None Vertex:{}"
                         .format(self.__vertex.get_vertex_id(),
                                 self.name,
                                 self.get_vertex().is_none()))

            if not self.get_vertex().is_none():
                # Initialize pyphot
                self.__execute_phot_init_function()
                self.__execute_normal_vertex_processing_logic()
            else:
                self.__execute_none_vertex_processing_logic()
        except Exception as e:
            logging.error("Arithmetic Unit Thread:{},vertex:{} got an exception:{}"
                          .format(self.name, self.__vertex.get_vertex_id(), e))
            self.__arithmetic_unit_status = ArithmeticUnitStatus.THREAD_EXCEPTION
            self.__error_queue.put(e)
